import os
import json
import zipfile
import tempfile
import shutil
import uuid
import subprocess
from keras.models import load_model
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import tensorflow as tf

app = FastAPI()
# 允许跨域访问
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 设置为 "*" 允许所有来源访问
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/convert_model/")
async def convert_model(file: UploadFile = File(...), model_type: str = Form('tflite')):
    """
    接收一个ZIP文件，并根据model_type参数决定返回Keras模型或TensorFlow Lite模型
    model_type: "keras" 或 "tflite"
    """
    # 创建临时目录
    temp_dir = tempfile.mkdtemp()

    # 保存上传的zip文件到临时目录
    dir_name = str(uuid.uuid4())
    zip_path = os.path.join(temp_dir, dir_name)
    with open(zip_path, "wb") as f:
        f.write(await file.read())

    # 解压zip文件
    try:
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(temp_dir)
    except zipfile.BadZipFile:
        message = '错误：不是zip压缩包'
        return JSONResponse(status_code=400, content={'message': message})

    # 解压后的文件路径
    tfjs_model_path = os.path.join(temp_dir, "model.json")
    metadata_path = os.path.join(temp_dir, "metadata.json")
    weight_path = os.path.join(temp_dir, "weights.bin")
    if not os.path.exists(tfjs_model_path) or not os.path.exists(metadata_path) or not os.path.exists(weight_path):
        shutil.rmtree(temp_dir)  # 清理临时目录
        message = '错误：缺少必要的文件，请导入正确的模型压缩包'
        return JSONResponse(status_code=400, content={'message': message})

    # 生成 labels.txt 文件
    with open(metadata_path, "r") as metadata_file:
        metadata = json.load(metadata_file)
    labels = metadata.get("labels", [])
    labels_txt_path = os.path.join(temp_dir, "labels.txt")
    with open(labels_txt_path, "w") as labels_file:
        for i, label in enumerate(labels):
            labels_file.write(f"{i} {label.strip()}\n")  # 去除空格并写入

    # 输出路径，用于保存转换后的模型
    output_dir = tempfile.mkdtemp()
    # 调用 tensorflowjs_converter 命令将 TensorFlow.js 模型转换为 Keras SavedModel
    keras_model_path = os.path.join(output_dir, "keras_model.h5")
    try:
        subprocess.run([
            "tensorflowjs_converter",
            "--input_format=tfjs_layers_model",
            "--output_format=keras",
            tfjs_model_path,
            keras_model_path
        ], check=True)
    except Exception as e:
        message = f'错误：{str(e)}'
        shutil.rmtree(temp_dir)  # 清理临时目录
        shutil.rmtree(output_dir)  # 清理输出目录
        return JSONResponse(status_code=400, content={'message': message})

    # 转换后的 Keras 模型路径
    if not os.path.exists(keras_model_path):
        shutil.rmtree(temp_dir)  # 清理临时目录
        shutil.rmtree(output_dir)  # 清理输出目录
        message = f'错误：keras 转换错误'
        return JSONResponse(status_code=500, content={'message': message})

    # 如果选择返回 Keras 模型
    if model_type == "keras":
        file_name = 'converted_keras.zip'
        output_zip = tempfile.mktemp(suffix='.zip')
        with zipfile.ZipFile(output_zip, 'w') as zipf:
            zipf.write(keras_model_path, "keras_model.h5")
            zipf.write(labels_txt_path, "labels.txt")

    # 如果选择返回 TensorFlow Lite 模型
    else:
        file_name = 'converted_tflite.zip'
        # 加载 Keras 模型
        keras_model = load_model(keras_model_path)

        # 转换为 TensorFlow Lite 模型
        converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
        tflite_model = converter.convert()

        # 保存为 tflite 模型
        tflite_model_path = os.path.join(output_dir, "keras_model.tflite")
        with open(tflite_model_path, "wb") as f:
            f.write(tflite_model)

        # 压缩为zip
        output_zip = tempfile.mktemp(suffix='.zip')
        with zipfile.ZipFile(output_zip, 'w') as zipf:
            zipf.write(tflite_model_path, "keras_model.tflite")
            zipf.write(labels_txt_path, "labels.txt")

    # 清理临时文件夹
    shutil.rmtree(temp_dir)
    shutil.rmtree(output_dir)

    return FileResponse(output_zip, media_type='application/zip',
                        headers={"Content-Disposition": f"attachment; filename={file_name}"})


# if __name__ == '__main__':
#     import uvicorn
#
#     uvicorn.run(app, host='0.0.0.0', port=18888)
